#!/usr/bin/python

import random

from hips.scheduling import schedule


class HipsConstraint(object):

  def __init__(self, doc_count, schedule, log):
    self._doc_count = doc_count
    self._schedule = schedule
    self._shift_count = schedule.ShiftCount()
    self._log = log

  def _CreateReason(self, description):
    return self._log.CreateConflictReason(description)

  def _LogConflict(self, reason, shift, doctor):
    self._log.LogConflict(reason, shift, doctor)

  def ForwardCheck(self, variables, domains, assignments):
    return True

  def PreProcess(self, variables, domains, constraints, vconstraints):
    pass


class NoTwoConsecutiveShifts(HipsConstraint):

  _FORCE_DETAILED_CHECK = True

  def __init__(self, doc_count, schedule, log):
    HipsConstraint.__init__(self, doc_count, schedule, log)
    self.two_consecutive_shifts = self._CreateReason("two consecutive shifts")

  def __call__(self, variables, domains, assignments, forwardcheck=False):
    # In theory, we coud skip detailed check here as forward checks
    # should eliminate conflicts. However this doesn't actually result
    # in performance gain.
    if not NoTwoConsecutiveShifts._FORCE_DETAILED_CHECK:
      return 0
    conflicts = 0
    for i in xrange(self._shift_count - 1):
      if variables[i] not in assignments or variables[i + 1] not in assignments:
        continue
      if assignments[variables[i]] & assignments[variables[i + 1]]:
        overlap = assignments[variables[i]] & assignments[variables[i + 1]]
        for j in range(self._doc_count):
          if overlap & (1 << j):
            self._LogConflict(self.two_consecutive_shifts, i, j)
            conflicts += 1
    if forwardcheck:
      return self.ForwardCheck(variables, domains, assignments)
    return conflicts

  def ForwardCheck(self, variables, domains, assignments):
    # Look at surrounding values in any unassigned variables to prune.
    for i, variable in enumerate(variables):
      # In a general case solver, a check for "not in assignments"
      # would make more sense, but the min conflict solver will
      # re-assign variables that are already assigned.
      # Instead, in order to speed this up, we fix the night shifts
      # to reduce our search space.
      if not self._schedule.IsDay(variable):
        continue
      surrounding_values = 0
      if i > 0 and variables[i - 1] in assignments:
        surrounding_values |= assignments[variables[i - 1]]
      if i < len(variables) - 1 and variables[i + 1] in assignments:
        surrounding_values |= assignments[variables[i + 1]]
      domain = domains[variable]
      for value in domain:
        if value & surrounding_values:
          domain.HideValue(value)
          if not domain:
            return False
    return True


def HideDoc(variables, domains, assignments, index, doc_mask, shift_count):
  if index >= 0 and index < shift_count and variables[index] not in assignments:
    domain = domains[index]
    for value in domain[:]:
      if value & doc_mask:
        domain.HideValue(value)
        if not domain:
          return False
  return True


class NoMoreThanMaxConsecutiveDays(HipsConstraint):

  def __init__(self, day_counts, doc_count, schedule, log):
    HipsConstraint.__init__(self, doc_count, schedule, log)
    self._day_counts = day_counts
    self.too_many_consecutive_shifts = self._CreateReason(
        "too many consecutive shifts")

  def __call__(self, variables, domains, assignments, forwardcheck=False):
    current_len_found = [0] * self._doc_count
    max_len_found = [0] * self._doc_count
    # TODO There could be more than one run.
    max_len_end_index = [-1] * self._doc_count
    conflicts = 0
    # TODO Use schedule API
    for i, variable in enumerate(variables):
      if not self._schedule.IsDay(variable):
        continue
      if variables[i] not in assignments or variables[i + 1] not in assignments:
        current_len_found = [0] * self._doc_count
        continue
      day = assignments[variables[i]] | assignments[variables[i + 1]]
      for j in xrange(self._doc_count):
        if day & (1 << j):
          current_len_found[j] += 1
          if current_len_found[j] > self._day_counts[j]:
            self._LogConflict(self.too_many_consecutive_shifts, variable, j)
            conflicts += 1
          if current_len_found[j] > max_len_found[j]:
            max_len_found[j] = current_len_found[j]
            max_len_end_index[j] = i
        else:
          current_len_found[j] = 0
    if forwardcheck:
      return self.ForwardCheck(variables, domains, assignments)
    return conflicts

  def ForwardCheck(self, variables, domains, assignments):
    # If we're at the maximum number of consecutive days,
    # disallow the doc for the next day.
    # for j in xrange(self._doc_count):
    #   if max_len_found[j] == self._day_counts[j]:
    #     doc_mask = 1 << j
    #     # Wrap to nearest beginning of day if necessary
    #     if max_len_end_index[j] & 1:
    #       max_len_end_index[j] -= 1
    #     # before_index is index at one day prior to shift
    #     before_index = max_len_end_index[j] - self._day_counts[j] * 2
    #     if (not HideDoc(variables, domains, assignments,
    #             before_index, doc_mask, self._shift_count) or 
    #         not HideDoc(variables, domains, assignments,
    #             before_index + 1, doc_mask, self._shift_count)):
    #       return False
    #     after_index = max_len_end_index[j] + 2
    #     if (not HideDoc(variables, domains, assignments,
    #             after_index, doc_mask, self._shift_count) or
    #         not HideDoc(variables, domains, assignments,
    #             after_index + 1, doc_mask, self._shift_count)):
    #       return False
    return True


class TotalShiftsInRange(HipsConstraint):

  def __init__(self, min_shift_counts, max_shift_counts,
      max_night_shift_counts, doc_count, schedule, log):
    HipsConstraint.__init__(self, doc_count, schedule, log)
    self._min_shift_counts = min_shift_counts
    self._max_shift_counts = max_shift_counts
    self._max_night_shift_counts = max_night_shift_counts
    self.too_many_shifts = self._CreateReason("too many shifts")
    self.too_many_night_shifts = self._CreateReason("too many night shifts")
    self.not_enough_shifts = self._CreateReason("not enough shifts")

  def __call__(self, variables, domains, assignments, forwardcheck=False):
    shift_counts = [0] * self._doc_count
    night_shift_counts = [0] * self._doc_count
    unassigned_shift_counts = [0] * self._doc_count
    too_many_shifts_conflicts = 0
    too_many_night_shifts_conflicts = 0
    not_enough_shifts_conflicts = 0
    for shift, value in assignments.iteritems():
      for j in xrange(self._doc_count):
        doc_mask = 1 << j
        if value & doc_mask:
          shift_counts[j] += 1
          if shift_counts[j] > self._max_shift_counts[j]:
            self._LogConflict(self.too_many_shifts, shift, j)
            too_many_shifts_conflicts += 1
          if self._schedule.IsNight(shift):
            night_shift_counts[j] += 1
            if night_shift_counts[j] > self._max_night_shift_counts[j]:
              self._LogConflict(self.too_many_night_shifts, shift, j)
              too_many_night_shifts_conflicts += 1
        else:
          unassigned_shift_counts[j] += 1
          if (self._shift_count - unassigned_shift_counts[j] <
              self._min_shift_counts[j]):
            self._LogConflict(self.not_enough_shifts, shift, j)
            not_enough_shifts_conflicts += 1
    if forwardcheck:
      return self.ForwardCheck(variables, domains, assignments)
    return (too_many_shifts_conflicts + too_many_night_shifts_conflicts +
        not_enough_shifts_conflicts)

  def ForwardCheck(self, variables, domains, assignments):
    # If we're at the maximum number of shifts, don't allow any more.
    # for j in xrange(self._doc_count):
    #   doc_mask = 1 << j
    #   if shift_counts[j] == self._max_shift_counts[j]:
    #     for i in xrange(self._shift_count):
    #       if not HideDoc(variables, domains, assignments, i, doc_mask,
    #           self._shift_count):
    #         return False
    return True


class NotInDay(HipsConstraint):

  def __init__(self, days, doc_mask, doc_count, schedule, log):
    HipsConstraint.__init__(self, doc_count, schedule, log)
    self._days = days
    self._doc_mask = doc_mask

  def PreProcess(self, variables, domains, constraints, vconstraints):
    for i, variable in enumerate(variables):
      domain = domains[variable]
      for value in domain[:]:
        if (self._schedule.DayForShiftNumber(i) in self._days and
            value & self._doc_mask):
          domain.remove(value)
      vconstraints[variable].remove((self, variables))
    constraints.remove((self, variables))

  def __call__(self, variables, domains, assignments, forwardcheck=False):
    # PreProcess() will remove it.
    raise RuntimeError, "Can't happen"


class NoLessThanMinConsecutiveDayShifts(HipsConstraint):

  def __init__(self, doc_count, schedule, log, min_len=3):
    HipsConstraint.__init__(self, doc_count, schedule, log)
    self.min_len = min_len
    self.too_few_consecutive_shifts = self._CreateReason(
        "too few consecutive shifts")
    
  def __call__(self, variables, domains, assignments, forwardcheck=False):
    current_len_found = [0] * self._doc_count
    conflicts = 0
    for i, variable in enumerate(variables):
      if not self._schedule.IsDay(variable):
        continue
      # Reset counts if we find an unassigned day.
      # Hmm, this shouldn't happen.
      if variables[i] not in assignments:
        current_len_found = [0] * self._doc_count
        continue
      day = assignments[variables[i]]
      for j in xrange(self._doc_count):
        if day & (1 << j):
          current_len_found[j] += 1
        else:
          if (current_len_found[j] > 0 and
              current_len_found[j] < self.min_len and
              day > 2):
            self._LogConflict(self.too_few_consecutive_shifts, variable, j)
            conflicts += 1
          current_len_found[j] = 0
    if forwardcheck:
      return self.ForwardCheck(variables, domains, assignments)
    return conflicts


class DayShiftsOverlap(HipsConstraint):

  def __init__(self, doc_count, schedule, log):
    HipsConstraint.__init__(self, doc_count, schedule, log)
    self.no_overlapping_doctors = self._CreateReason("no overlapping doctors")

  def __call__(self, variables, domains, assignments, forwardcheck=False):
    conflicts = 0
    for i, variable in enumerate(variables):
      if self._schedule.IsDay(variable) and i > 0:
        # Reset counts if we find an unassigned day.
        # Hmm, this shouldn't happen.
        if variables[i] not in assignments:
          continue
        day = assignments[variables[i]]
        prev_day = assignments[variables[i - 2]]
        if day & prev_day == 0:
          self._LogConflict(self.no_overlapping_doctors, variable, 0)
          conflicts += 1
    if forwardcheck:
      return self.ForwardCheck(variables, domains, assignments)
    return conflicts
